{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fee27f39",
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "from openai import OpenAI\n",
    "import gradio as gr\n",
    "\n",
    "load_dotenv(override=True)\n",
    "openai_api_key = os.getenv('OPENAI_API_KEY')\n",
    "anthropic_api_key = os.getenv('ANTHROPIC_API_KEY')\n",
    "google_api_key = os.getenv('GOOGLE_API_KEY')\n",
    "ollama_api_key = os.getenv('OLLAMA_API_KEY')\n",
    "\n",
    "if openai_api_key:\n",
    "    print(f\"OpenAI API Key exists and begins {openai_api_key[:8]}\")\n",
    "else:\n",
    "    print(\"OpenAI API Key not set\")\n",
    "    \n",
    "if anthropic_api_key:\n",
    "    print(f\"Anthropic API Key exists and begins {anthropic_api_key[:7]}\")\n",
    "else:\n",
    "    print(\"Anthropic API Key not set (and this is optional)\")\n",
    "\n",
    "if google_api_key:\n",
    "    print(f\"Google API Key exists and begins {google_api_key[:2]}\")\n",
    "else:\n",
    "    print(\"Google API Key not set (and this is optional)\")\n",
    "\n",
    "if ollama_api_key:\n",
    "    print(f\"OLLAMA API Key exists and begins {ollama_api_key[:2]}\")\n",
    "else:\n",
    "    print(\"OLLAMA API Key not set (and this is optional)\")\n",
    "\n",
    "# Connect to client libraries\n",
    "\n",
    "openai = OpenAI()\n",
    "\n",
    "anthropic_url = \"https://api.anthropic.com/v1/\"\n",
    "gemini_url = \"https://generativelanguage.googleapis.com/v1beta/openai/\"\n",
    "ollama_url = \"http://localhost:11434/v1\"\n",
    "\n",
    "anthropic = OpenAI(api_key=anthropic_api_key, base_url=anthropic_url)\n",
    "gemini = OpenAI(api_key=google_api_key, base_url=gemini_url)\n",
    "ollama = OpenAI(api_key=ollama_api_key, base_url=ollama_url)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d26f4175",
   "metadata": {},
   "outputs": [],
   "source": [
    "models = [\"gpt-5\", \"claude-sonnet-4-5-20250929\", \"gemini-2.5-pro\", \"gpt-oss:20b-cloud\", ]\n",
    "\n",
    "clients = {\"gpt-5\": openai, \"claude-sonnet-4-5-20250929\": anthropic, \"gemini-2.5-pro\": gemini,  \"gpt-oss:20b-cloud\": ollama}\n",
    "\n",
    "# Want to keep costs ultra-low? Replace this with models of your choice, using the examples from yesterday"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76563884",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt_doc = \"\"\"You are an expert Python developer and code reviewer.\n",
    "Your job is to read the user's provided function, and return:\n",
    "1. A concise, PEP-257-compliant docstring summarizing what the function does, clarifying types, parameters, return values, and side effects.\n",
    "2. Helpful inline comments that improve both readability and maintainability, without restating what the code obviously does.\n",
    "\n",
    "Only output the function, not explanations or additional text. \n",
    "Do not modify variable names or refactor the function logic.\n",
    "Your response should improve the code's clarity and documentation, making it easier for others to understand and maintain.\n",
    "Don't be extremely verbose.\n",
    "Your answer should be at a {level} level of expertise.\n",
    "\"\"\"\n",
    "\n",
    "system_prompt_tests = \"\"\"You are a seasoned Python developer and testing expert.\n",
    "Your task is to read the user's provided function, and generate:\n",
    "1. A concise set of meaningful unit tests that thoroughly validate the function's correctness, including typical, edge, and error cases.\n",
    "2. The tests should be written for pytest (or unittest if pytest is not appropriate), use clear, descriptive names, and avoid unnecessary complexity.\n",
    "3. If dependencies or mocking are needed, include minimal necessary setup code (but avoid over-mocking).\n",
    "\n",
    "Only output the relevant test code, not explanations or extra text.\n",
    "Do not change the original function; focus solely on comprehensive, maintainable test coverage that other developers can easily understand and extend.\n",
    "\"\"\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bd82e96",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_documentation(code, model, level):\n",
    "    response = clients[model].chat.completions.create(\n",
    "        model=model,\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": system_prompt_doc.format(level=level)},\n",
    "            {\"role\": \"user\", \"content\": code}\n",
    "        ],\n",
    "        stream=True\n",
    "    )\n",
    "    output = \"\"\n",
    "    for chunk in response:\n",
    "        output += chunk.choices[0].delta.content or \"\"\n",
    "        yield output.replace(\"```python\", \"\").replace(\"```\", \"\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b01b3421",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_tests(code, model ):\n",
    "    response = clients[model].chat.completions.create(\n",
    "        model=model,\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": system_prompt_tests},\n",
    "            {\"role\": \"user\", \"content\": code}\n",
    "        ],\n",
    "        stream=True\n",
    "    )\n",
    "    output = \"\"\n",
    "    for chunk in response:\n",
    "        output += chunk.choices[0].delta.content or \"\"\n",
    "        yield output.replace(\"```python\", \"\").replace(\"```\", \"\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16b71915",
   "metadata": {},
   "outputs": [],
   "source": [
    "vscode_dark = gr.themes.Monochrome(\n",
    "    primary_hue=\"blue\",\n",
    "    secondary_hue=\"slate\",\n",
    "    neutral_hue=\"slate\",\n",
    ").set(\n",
    "    body_background_fill=\"#1e1e1e\",\n",
    "    body_background_fill_dark=\"#1e1e1e\",\n",
    "    block_background_fill=\"#252526\",\n",
    "    block_background_fill_dark=\"#252526\",\n",
    "    block_border_color=\"#3e3e42\",\n",
    "    block_border_color_dark=\"#3e3e42\",\n",
    "    border_color_primary=\"#3e3e42\",\n",
    "    block_label_background_fill=\"#252526\",\n",
    "    block_label_background_fill_dark=\"#252526\",\n",
    "    block_label_text_color=\"#cccccc\",\n",
    "    block_label_text_color_dark=\"#cccccc\",\n",
    "    block_title_text_color=\"#cccccc\",\n",
    "    block_title_text_color_dark=\"#cccccc\",\n",
    "    body_text_color=\"#d4d4d4\",\n",
    "    body_text_color_dark=\"#d4d4d4\",\n",
    "    button_primary_background_fill=\"#0e639c\",\n",
    "    button_primary_background_fill_dark=\"#0e639c\",\n",
    "    button_primary_background_fill_hover=\"#1177bb\",\n",
    "    button_primary_background_fill_hover_dark=\"#1177bb\",\n",
    "    button_primary_text_color=\"#ffffff\",\n",
    "    button_primary_text_color_dark=\"#ffffff\",\n",
    "    input_background_fill=\"#3c3c3c\",\n",
    "    input_background_fill_dark=\"#3c3c3c\",\n",
    "    color_accent=\"#007acc\",\n",
    "    color_accent_soft=\"#094771\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23311022",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "\n",
    "with gr.Blocks(theme=vscode_dark, css=\"\"\"\n",
    "    .gradio-container {font-family: 'Consolas', 'Monaco', monospace;}\n",
    "    h1 {color: #d4d4d4 !important;}\n",
    "\"\"\") as ui:\n",
    "    gr.Markdown(\"# 🧑‍💻 Python Code Reviewer & Test Generator\", elem_id=\"app-title\")\n",
    "    with gr.Tab(\"Docstring & Comments\") as tab1:\n",
    "        gr.Markdown(\"# Function Docstring & Comment Helper\\nPaste your function below and get helpful docstrings and inline comments!\")\n",
    "\n",
    "        with gr.Row():\n",
    "            code_input_1 = gr.Code(label=\"Paste your Python function here\", lines=10, language=\"python\")\n",
    "            code_output = gr.Code(label=\"Function with improved docstring and comments\", lines=10, language=\"python\")\n",
    "        \n",
    "        with gr.Row(equal_height=True):\n",
    "            level_radio = gr.Radio(choices=[\"Junior\", \"Mid\", \"Senior\"], value=\"Mid\", label=\"Reviewer level\", interactive=True)\n",
    "            model_dropdown = gr.Dropdown(choices=models, value=models[-1], label=\"Select model\")\n",
    "            submit_doc_btn = gr.Button(\"Generate docstring & comments\", scale=0.5)\n",
    "\n",
    "        submit_doc_btn.click(\n",
    "            generate_documentation, \n",
    "            inputs=[code_input_1, model_dropdown, level_radio], \n",
    "            outputs=code_output\n",
    "        )\n",
    "\n",
    "    with gr.Tab(\"Unit Tests\") as tab2:\n",
    "        gr.Markdown(\"# Unit Test Generator\\nPaste your function below and get auto-generated unit tests!\")\n",
    "\n",
    "        with gr.Row():\n",
    "            code_input_2 = gr.Code(label=\"Paste your Python function here\", lines=10, language=\"python\")\n",
    "            code_output_2 = gr.Code(label=\"Generated tests\", lines=10, language=\"python\")\n",
    "        \n",
    "        with gr.Row(equal_height=True):\n",
    "            model_dropdown_2 = gr.Dropdown(choices=models, value=models[-1], label=\"Select model\")\n",
    "            submit_test_btn = gr.Button(\"Generate unit tests\", scale=0.5)\n",
    "\n",
    "        submit_test_btn.click(\n",
    "            generate_tests, \n",
    "            inputs=[code_input_2, model_dropdown_2], \n",
    "            outputs=code_output_2\n",
    "        )\n",
    "    \n",
    "    tab2.select(lambda x: x, inputs=code_input_1, outputs=code_input_2)\n",
    "    tab1.select(lambda x: x, inputs=code_input_2, outputs=code_input_1)\n",
    "\n",
    "ui.launch(share=False, inbrowser=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
